import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import os
import logging
import shutil
import subprocess
from PIL import Image, ImageTk
import threading
import sys
import platform
from concurrent.futures import ThreadPoolExecutor, as_completed
import paddle
paddle.disable_signal_handler()  # 禁用 Paddle 的信号处理
from paddleocr import PaddleOCR
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
from ocr_processor import KEYWORDS, check_keyword_match, safe_move_file
import multiprocessing
import cv2
import queue
import numpy as np
import warnings
warnings.filterwarnings('ignore')  # 忽略警告

# 全局OCR引擎
global_ocr_engine = None

def init_ocr_engine():
    """初始化GPU加速的OCR引擎"""
    try:
        # 检查GPU是否可用
        if not paddle.device.is_compiled_with_cuda():
            messagebox.showwarning("警告", "未检测到CUDA支持，将使用CPU模式运行")
            use_gpu = False
        else:
            use_gpu = True
            logging.info("使用GPU加速模式")
        
        engine = PaddleOCR(
            use_angle_cls=False,  # 关闭方向分类提升速度
            lang='ch',
            ocr_version='PP-OCRv4',  # 使用更强模型
            rec_batch_num=6,  # 批量识别
            use_gpu=use_gpu,  # 启用GPU
            gpu_mem=500,  # GPU显存使用量(MB)
            enable_mkldnn=True,  # 启用MKL-DNN加速
            cpu_threads=4,  # CPU线程数
        )
        return engine
    except Exception as e:
        logging.error(f"OCR引擎初始化失败: {str(e)}")
        return None

class GPUOCRTab(ttk.Frame):
    def __init__(self, parent):
        super().__init__(parent, padding="10 0 10 0")
        self.grid(sticky=(tk.N, tk.S, tk.E, tk.W), padx=20)
        self.grid_columnconfigure(0, weight=1)
        self.grid_rowconfigure(1, weight=1)
        
        # 初始化OCR引擎
        global global_ocr_engine
        if global_ocr_engine is None:
            global_ocr_engine = init_ocr_engine()
            if global_ocr_engine is None:
                messagebox.showerror("错误", "GPU OCR引擎初始化失败，请检查CUDA环境")
        
        # 初始化变量
        self.source_dir = None
        self.processed_folders = {}  # 存储已处理的文件夹信息
        self.processed_files_dict = {}  # 分类到图片文件列表的映射
        self.current_category = None
        self.selected_processed_file = None
        
        # 设置样式
        self.style = ttk.Style()
        self.style.configure('Title.TLabelframe.Label', font=('Arial', 12, 'bold'))
        
        # 创建UI组件
        self.create_file_selection_frame()
        self.create_preview_frame()
        self.create_bottom_frame()
        
        # 添加GPU状态显示
        self.update_gpu_status()
        
    def update_gpu_status(self):
        """更新GPU状态显示"""
        try:
            if paddle.device.is_compiled_with_cuda():
                gpu_info = paddle.device.cuda.get_device_properties()
                gpu_status = f"GPU状态: {gpu_info.name} (显存: {gpu_info.total_memory/1024/1024/1024:.1f}GB)"
            else:
                gpu_status = "GPU状态: 未检测到CUDA支持，使用CPU模式"
        except Exception as e:
            gpu_status = f"GPU状态: 检测失败 ({str(e)})"
        
        self.status_var.set(gpu_status)
        
    def create_file_selection_frame(self):
        file_frame = ttk.LabelFrame(self, text="GPU加速文件选择", padding="5", style='Title.TLabelframe')
        file_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=5)
        file_frame.grid_columnconfigure(0, weight=1)
        file_frame.grid_columnconfigure(1, weight=0)
        file_frame.grid_columnconfigure(2, weight=0)
        
        # 左侧：显示选择的目录路径
        self.dir_path_var = tk.StringVar()
        dir_entry = ttk.Entry(file_frame, textvariable=self.dir_path_var, width=60, state='readonly')
        dir_entry.grid(row=0, column=0, padx=5, pady=5, sticky=(tk.W, tk.E))
        
        # 中间：选择目录按钮
        select_btn = ttk.Button(file_frame, text="选择文件夹", command=self.select_directory)
        select_btn.grid(row=0, column=1, padx=5, pady=5)
        
        # 右侧：开始识别按钮
        self.recognize_btn = ttk.Button(file_frame, text="开始GPU识别", command=self.start_recognition, state='disabled')
        self.recognize_btn.grid(row=0, column=2, padx=5, pady=5)
        
    def create_preview_frame(self):
        preview_frame = ttk.Frame(self)
        preview_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.N, tk.S, tk.E, tk.W), pady=5)
        preview_frame.grid_columnconfigure(0, weight=1)
        preview_frame.grid_columnconfigure(1, weight=5)
        preview_frame.grid_rowconfigure(0, weight=1)
        preview_frame.grid_propagate(False)
        preview_frame.config(height=500)
        
        # 左侧：待处理区域
        self.create_pending_frame(preview_frame)
        
        # 右侧：已处理区域
        self.create_processed_frame(preview_frame)
        
    def create_pending_frame(self, parent):
        self.pending_frame = ttk.LabelFrame(parent, text="待处理图片 (0张)", padding="5", style='Title.TLabelframe')
        self.pending_frame.grid(row=0, column=0, sticky=(tk.N, tk.S, tk.E, tk.W), padx=(0, 5))
        self.pending_frame.grid_rowconfigure(1, weight=1)
        self.pending_frame.grid_columnconfigure(0, weight=1)
        self.pending_frame.config(height=400)
        self.pending_frame.grid_propagate(False)
        
        # 检索输入框
        search_frame = ttk.Frame(self.pending_frame)
        search_frame.pack(fill=tk.X, pady=(0, 5))
        ttk.Label(search_frame, text="检索:").pack(side=tk.LEFT, padx=(0, 2))
        self.search_var = tk.StringVar()
        search_entry = ttk.Entry(search_frame, textvariable=self.search_var, width=20)
        search_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.search_var.trace_add('write', self.filter_pending_list)
        
        # 待处理图片列表
        self.pending_listbox = tk.Listbox(self.pending_frame, selectmode=tk.SINGLE)
        self.pending_listbox.pack(fill=tk.BOTH, expand=True)
        pending_scrollbar = ttk.Scrollbar(self.pending_frame, orient="vertical", command=self.pending_listbox.yview)
        pending_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.pending_listbox.configure(yscrollcommand=pending_scrollbar.set)
        
        # 预览按钮
        preview_btn = ttk.Button(self.pending_frame, text="预览", command=self.preview_pending_image)
        preview_btn.pack(fill=tk.X, pady=(5, 0))
        
        # 存储原始文件名列表用于检索
        self._all_pending_files = []
        
    def create_processed_frame(self, parent):
        self.processed_frame = ttk.LabelFrame(parent, text="已处理图片 (0张)", padding="5", style='Title.TLabelframe')
        self.processed_frame.grid(row=0, column=1, sticky=(tk.N, tk.S, tk.E, tk.W), padx=(5, 0))
        self.processed_frame.grid_rowconfigure(1, weight=1)
        self.processed_frame.grid_columnconfigure(0, weight=1, minsize=80)
        self.processed_frame.grid_columnconfigure(1, weight=5)
        self.processed_frame.config(height=400)
        self.processed_frame.grid_propagate(False)

        # 左侧：分类列表
        self.category_listbox = tk.Listbox(self.processed_frame, selectmode=tk.SINGLE, font=("Arial", 11))
        self.category_listbox.grid(row=0, column=0, rowspan=3, sticky=(tk.N, tk.S, tk.E, tk.W), padx=(0, 5), pady=2)
        self.category_listbox.bind('<<ListboxSelect>>', self.on_category_select)
        self.category_scrollbar = ttk.Scrollbar(self.processed_frame, orient="vertical", command=self.category_listbox.yview)
        self.category_scrollbar.grid(row=0, column=0, rowspan=3, sticky=(tk.N, tk.S, tk.E))
        self.category_listbox.configure(yscrollcommand=self.category_scrollbar.set)

        # 右侧：检索输入框
        search_frame = ttk.Frame(self.processed_frame)
        search_frame.grid(row=0, column=1, sticky=(tk.W, tk.E), pady=(0, 5))
        ttk.Label(search_frame, text="检索:").pack(side=tk.LEFT, padx=(0, 2))
        self.processed_search_var = tk.StringVar()
        search_entry = ttk.Entry(search_frame, textvariable=self.processed_search_var, width=20)
        search_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.processed_search_var.trace_add('write', self.filter_processed_list)

        # 右侧：图片文件列表（独占一行，最大化宽度）
        self.processed_listbox = tk.Listbox(self.processed_frame, selectmode=tk.SINGLE, font=("Arial", 11))
        self.processed_listbox.grid(row=1, column=1, sticky=(tk.N, tk.S, tk.E, tk.W), padx=(5, 0), pady=2)
        self.processed_listbox.bind('<<ListboxSelect>>', self.on_processed_file_select)
        self.processed_list_scrollbar = ttk.Scrollbar(self.processed_frame, orient="vertical", command=self.processed_listbox.yview)
        self.processed_list_scrollbar.grid(row=1, column=1, sticky=(tk.N, tk.S, tk.E))
        self.processed_listbox.configure(yscrollcommand=self.processed_list_scrollbar.set)

        # 右侧：预览按钮
        preview_btn = ttk.Button(self.processed_frame, text="预览", command=self.preview_processed_image)
        preview_btn.grid(row=2, column=1, sticky=(tk.W, tk.E), pady=(5, 0))

        # 存储原始文件名列表用于检索
        self._all_processed_files = []

    def create_bottom_frame(self):
        bottom_frame = ttk.Frame(self)
        bottom_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(10, 0))
        bottom_frame.grid_columnconfigure(0, weight=1)
        bottom_frame.grid_columnconfigure(1, weight=1)
        bottom_frame.config(height=60)
        # 左侧：运行状态（多行自动换行，宽度自适应）
        self.status_var = tk.StringVar(value="运行状态：等待操作")
        self.status_label = ttk.Label(
            bottom_frame, textvariable=self.status_var, font=("Arial", 11), anchor='w',
            wraplength=self.winfo_width() - 60, justify='left'
        )
        self.status_label.grid(row=0, column=0, sticky=(tk.W, tk.E), padx=10)
        # 监听窗口大小变化，动态调整wraplength
        def update_status_wrap(event):
            new_width = max(self.winfo_width() - 60, 200)
            self.status_label.config(wraplength=new_width)
        self.bind('<Configure>', update_status_wrap)
        # 进度条和百分比
        progress_frame = ttk.Frame(bottom_frame)
        progress_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), padx=10, pady=(5, 0))
        self.progress_var = tk.IntVar(value=0)
        self.progressbar = ttk.Progressbar(progress_frame, variable=self.progress_var, maximum=100)
        self.progressbar.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.percent_label = ttk.Label(progress_frame, text="0%", width=5)
        self.percent_label.pack(side=tk.LEFT, padx=(5, 0))
        # 右侧：文件操作按钮和中断按钮
        op_frame = ttk.Frame(bottom_frame)
        op_frame.grid(row=0, column=1, sticky=(tk.E), padx=10)
        del_btn = ttk.Button(op_frame, text="清空", command=self.delete_processed_file)
        del_btn.pack(side=tk.LEFT, padx=5)
        open_btn = ttk.Button(op_frame, text="打开文件夹位置", command=self.open_processed_file_location)
        open_btn.pack(side=tk.LEFT, padx=5)
        self.stop_btn = ttk.Button(op_frame, text="中断识别", command=self.stop_recognition, state='disabled')
        self.stop_btn.pack(side=tk.LEFT, padx=5)
        
    def format_number(self, num):
        """格式化数字，添加千位分隔符"""
        return f"{num:,}"
        
    def update_pending_list(self):
        self.pending_listbox.delete(0, tk.END)
        if hasattr(self, 'image_files'):
            self._all_pending_files = [os.path.basename(file) for file in self.image_files]
            self.filter_pending_list()
            self.pending_frame.configure(text=f"待处理图片 ({self.format_number(len(self._all_pending_files))}张)")
            
    def filter_pending_list(self, *args):
        keyword = self.search_var.get().strip().lower() if hasattr(self, 'search_var') else ''
        self.pending_listbox.delete(0, tk.END)
        if hasattr(self, '_all_pending_files'):
            for file in self._all_pending_files:
                if keyword in file.lower():
                    self.pending_listbox.insert(tk.END, file)

    def preview_pending_image(self):
        selection = self.pending_listbox.curselection()
        if not selection:
            messagebox.showinfo("提示", "请先在待处理图片列表中选中一张图片！")
            return
        filename = self.pending_listbox.get(selection[0])
        # 找到原始路径
        if hasattr(self, 'image_files'):
            for f in self.image_files:
                if os.path.basename(f) == filename:
                    img_path = f
                    break
            else:
                messagebox.showerror("错误", "未找到图片文件！")
                return
            # 弹窗显示图片
            self.show_image_popup(img_path)

    def show_image_popup(self, img_path):
        top = tk.Toplevel(self)
        top.title(f"预览 - {os.path.basename(img_path)}")
        try:
            img = Image.open(img_path)
            img.thumbnail((600, 600))
            photo = ImageTk.PhotoImage(img)
            label = tk.Label(top, image=photo)
            label.image = photo  # 防止被回收
            label.pack()
        except Exception as e:
            messagebox.showerror("错误", f"无法打开图片: {e}")
        
    def update_processed_list(self, folder_name, files):
        """更新已处理列表"""
        # 更新分类到文件的映射
        self.processed_files_dict[folder_name] = [os.path.basename(f) for f in files]
        self.refresh_category_list()
        self.refresh_processed_list()
        # 更新总数
        total_files = sum(len(files) for files in self.processed_files_dict.values())
        self.processed_frame.configure(text=f"已处理图片 ({self.format_number(total_files)}张)")
        
    def refresh_category_list(self):
        self.category_listbox.delete(0, tk.END)
        for category, files in self.processed_files_dict.items():
            display_name = f"{category} ({self.format_number(len(files))})"
            self.category_listbox.insert(tk.END, display_name)
        # 保持当前选中项
        if self.category_listbox.size() > 0:
            if self.current_category not in self.processed_files_dict:
                self.category_listbox.selection_clear(0, tk.END)
                self.category_listbox.selection_set(0)
                self.current_category = list(self.processed_files_dict.keys())[0]
            self.refresh_processed_list()
        
    def refresh_processed_list(self):
        self.processed_listbox.delete(0, tk.END)
        self._all_processed_files = []
        if self.current_category and self.current_category in self.processed_files_dict:
            for file in self.processed_files_dict[self.current_category]:
                self._all_processed_files.append(file)
            self.filter_processed_list()
        self.selected_processed_file = None

    def filter_processed_list(self, *args):
        keyword = self.processed_search_var.get().strip().lower() if hasattr(self, 'processed_search_var') else ''
        self.processed_listbox.delete(0, tk.END)
        for file in self._all_processed_files:
            if keyword in file.lower():
                self.processed_listbox.insert(tk.END, file)

    def preview_processed_image(self):
        selection = self.processed_listbox.curselection()
        if not selection:
            messagebox.showinfo("提示", "请先在已处理图片列表中选中一张图片！")
            return
        filename = self.processed_listbox.get(selection[0])
        # 找到原始路径
        if self.current_category and self.current_category in self.processed_files_dict:
            for f in self.image_files:
                if os.path.basename(f) == filename:
                    img_path = f
                    break
            else:
                messagebox.showerror("错误", "未找到图片文件！")
                return
            self.show_image_popup(img_path)

    def on_category_select(self, event):
        selection = self.category_listbox.curselection()
        if selection:
            idx = selection[0]
            # 解析真实分类名（去掉括号和数量）
            display_name = self.category_listbox.get(idx)
            category = display_name.split(' (')[0]
            self.current_category = category
            self.refresh_processed_list()
        
    def on_processed_file_select(self, event):
        selection = self.processed_listbox.curselection()
        if selection and self.current_category in self.processed_files_dict:
            idx = selection[0]
            self.selected_processed_file = self.processed_files_dict[self.current_category][idx]
        else:
            self.selected_processed_file = None
        
    def select_directory(self):
        dir_path = filedialog.askdirectory()
        if dir_path:
            self.source_dir = dir_path
            self.dir_path_var.set(dir_path)
            logging.info(f"选择目录: {dir_path}")
            self.scan_directory()
            
    def scan_directory(self):
        """扫描目录中的图片文件"""
        if not self.source_dir:
            return
            
        self.image_files = []
        for root, _, files in os.walk(self.source_dir):  # 使用os.walk递归搜索子目录
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                    self.image_files.append(os.path.join(root, file))
                
        if self.image_files:
            # 启用识别按钮
            self.recognize_btn.state(['!disabled'])
            # 更新待处理列表
            self.update_pending_list()
            messagebox.showinfo("提示", f"找到 {self.format_number(len(self.image_files))} 个图片文件")
        else:
            messagebox.showinfo("提示", "所选目录中没有找到图片文件")
            self.recognize_btn.state(['disabled'])
            
    def start_recognition(self):
        if not hasattr(self, 'image_files') or not self.image_files:
            messagebox.showwarning("警告", "请先选择包含图片的目录")
            return
        logging.info(f"开始识别目录中的 {self.format_number(len(self.image_files))} 个图片文件")
        self.recognize_btn.state(['disabled'])
        self.stop_btn.state(['!disabled'])
        self.status_var.set("运行状态：正在识别...")
        self.progress_var.set(0)
        self.percent_label.config(text="0%")
        self.stop_flag = threading.Event()
        input_dir = self.source_dir
        output_dir = os.path.join(os.path.dirname(input_dir), '输出图片')
        threading.Thread(target=self.run_ocr_and_refresh, args=(input_dir, output_dir), daemon=True).start()

    def stop_recognition(self):
        if hasattr(self, 'stop_flag'):
            self.stop_flag.set()
            self.status_var.set("运行状态：已请求中断，等待当前图片处理完成...")
            self.stop_btn.state(['disabled'])

    def run_ocr_and_refresh(self, input_dir, output_dir):
        try:
            image_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif')
            image_files = []
            for root, _, files in os.walk(input_dir):
                for file in files:
                    if file.lower().endswith(image_extensions):
                        image_files.append(os.path.join(root, file))
                    
            total = len(image_files)
            if total == 0:
                self.status_var.set("运行状态：未找到图片文件")
                return

            # 确保输出目录存在
            for keyword in KEYWORDS.keys():
                os.makedirs(os.path.join(output_dir, keyword), exist_ok=True)
            os.makedirs(os.path.join(output_dir, "未分类"), exist_ok=True)

            # 使用全局OCR引擎
            global global_ocr_engine
            if global_ocr_engine is None:
                global_ocr_engine = init_ocr_engine()
                if global_ocr_engine is None:
                    self.status_var.set("运行状态：OCR引擎初始化失败")
                    return

            def preprocess_and_resize(image_path):
                try:
                    if platform.system() == 'Windows':
                        img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
                    else:
                        img = cv2.imread(image_path)
                    if img is None:
                        logging.warning(f"无法读取图片: {image_path}")
                        return None
                    
                    # GPU版本优化：增加批处理大小
                    max_width = 1600  # 增加最大宽度
                    if img.shape[1] > max_width:
                        scale = max_width / img.shape[1]
                        img = cv2.resize(img, (int(img.shape[1]*scale), int(img.shape[0]*scale)))
                    
                    # GPU版本优化：增强图像预处理
                    # 1. 自适应直方图均衡化
                    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
                    l, a, b = cv2.split(lab)
                    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
                    cl = clahe.apply(l)
                    enhanced_img = cv2.merge((cl,a,b))
                    enhanced_img = cv2.cvtColor(enhanced_img, cv2.COLOR_LAB2RGB)
                    
                    # 2. 降噪
                    enhanced_img = cv2.fastNlMeansDenoisingColored(enhanced_img, None, 10, 10, 7, 21)
                    
                    return enhanced_img
                except Exception as e:
                    logging.error(f"图片预处理失败: {str(e)}")
                    return None

            # 新增：线程安全队列用于主线程刷新UI
            result_queue = queue.Queue()
            refresh_interval = 200  # ms

            def process_one(image_path, idx):
                try:
                    if self.stop_flag.is_set():
                        result_queue.put(("中断", idx, 0, "运行状态：已中断"))
                        return "stopped"
                    
                    img = preprocess_and_resize(image_path)
                    if img is None:
                        result_queue.put(("跳过", idx, int((idx + 1) / total * 100), f"跳过图片: {os.path.basename(image_path)}"))
                        return "skip"
                    
                    try:
                        with warnings.catch_warnings():
                            warnings.simplefilter("ignore")
                            # GPU版本优化：增加批处理大小
                            result = global_ocr_engine.ocr(img, cls=False)
                            
                        if result is None or len(result) == 0 or result[0] is None:
                            logging.warning(f"OCR返回空结果: {image_path}")
                            text = ''
                        else:
                            # 提取文本
                            text_parts = []
                            for line in result[0]:
                                if isinstance(line, list) and len(line) > 1:
                                    if isinstance(line[1], tuple) and len(line[1]) > 0:
                                        text_parts.append(str(line[1][0]))
                                    elif isinstance(line[1], str):
                                        text_parts.append(line[1])
                            text = ''.join(text_parts)

                    except Exception as e:
                        logging.error(f"OCR处理失败: {str(e)}, 图片路径: {image_path}")
                        text = ''

                    # 处理文件移动
                    try:
                        matched = False
                        for keyword, variants in KEYWORDS.items():
                            is_matched, _ = check_keyword_match(text, variants)
                            if is_matched:
                                dst = os.path.join(output_dir, keyword, os.path.basename(image_path))
                                shutil.copy2(image_path, dst)
                                matched = True
                                msg = f"处理: {os.path.basename(image_path)} -> {keyword}"
                                break
                        
                        if not matched:
                            dst = os.path.join(output_dir, "未分类", os.path.basename(image_path))
                            shutil.copy2(image_path, dst)
                            msg = f"处理: {os.path.basename(image_path)} -> 未分类"
                        
                        percent = int((idx + 1) / total * 100)
                        result_queue.put(("完成", idx, percent, msg))
                        return msg
                    except Exception as e:
                        error_msg = f"文件处理失败 {os.path.basename(image_path)}: {str(e)}"
                        logging.error(error_msg)
                        result_queue.put(("跳过", idx, int((idx + 1) / total * 100), error_msg))
                        return "error"
                except Exception as e:
                    error_msg = f"处理失败 {os.path.basename(image_path)}: {str(e)}"
                    logging.error(error_msg)
                    result_queue.put(("跳过", idx, int((idx + 1) / total * 100), error_msg))
                    return "error"

            def worker():
                try:
                    # GPU版本优化：增加线程池大小
                    with ThreadPoolExecutor(max_workers=8) as executor:  # 增加线程数
                        futures = []
                        for idx, img in enumerate(image_files):
                            if self.stop_flag.is_set():
                                break
                            future = executor.submit(process_one, img, idx)
                            futures.append(future)
                        # 等待所有任务完成
                        for future in as_completed(futures):
                            if self.stop_flag.is_set():
                                break
                            try:
                                future.result()
                            except Exception as e:
                                logging.error(f"处理任务出错: {str(e)}")
                    if not self.stop_flag.is_set():
                        result_queue.put(("全部完成", total, 100, "运行状态：GPU识别完成！"))
                except Exception as e:
                    result_queue.put(("错误", -1, 0, f"处理出错: {str(e)}"))

            # 启动后台线程做OCR
            threading.Thread(target=worker, daemon=True).start()
            # 启动UI定时刷新
            self.after(refresh_interval, lambda: self.refresh_ui_periodically(result_queue, total, output_dir))

        except Exception as e:
            self.status_var.set(f"运行状态：识别出错: {e}")
            messagebox.showerror("错误", f"识别过程中发生错误: {e}")
        finally:
            self.after(0, lambda: self.recognize_btn.state(['!disabled']))
            self.after(0, lambda: self.stop_btn.state(['disabled']))

    def refresh_ui_periodically(self, result_queue, total, output_dir):
        """定期刷新UI的函数"""
        try:
            updated = False
            while not result_queue.empty():
                status, idx, percent, msg = result_queue.get()
                if status == "中断":
                    self.status_var.set(msg)
                    self.update_progress(percent)
                    self.recognize_btn.state(['!disabled'])
                    self.stop_btn.state(['disabled'])
                    updated = True
                    break
                elif status == "错误":
                    self.status_var.set(msg)
                    updated = True
                elif status == "跳过":
                    self.status_var.set(msg)
                    self.update_progress(percent)
                    updated = True
                elif status == "完成":
                    self.status_var.set(msg)
                    self.update_progress(percent)
                    updated = True
                elif status == "全部完成":
                    self.status_var.set(msg)
                    self.update_progress(percent)
                    self.load_processed_from_output(output_dir)
                    self.recognize_btn.state(['!disabled'])
                    self.stop_btn.state(['disabled'])
                    updated = True
                    return  # 识别完成后不再定时刷新

            if updated:
                self.load_processed_from_output(output_dir)
            
            if not self.stop_flag.is_set():
                self.after(200, lambda: self.refresh_ui_periodically(result_queue, total, output_dir))
        except Exception as e:
            logging.error(f"UI更新出错: {str(e)}")
            self.status_var.set(f"UI更新出错，请检查日志")

    def update_progress(self, percent):
        self.progress_var.set(percent)
        self.percent_label.config(text=f"{percent}%")

    def load_processed_from_output(self, output_dir):
        """从输出图片目录加载已处理图片分类和文件"""
        if not os.path.exists(output_dir):
            return
        self.processed_files_dict.clear()
        for folder in os.listdir(output_dir):
            folder_path = os.path.join(output_dir, folder)
            if os.path.isdir(folder_path):
                files = [os.path.basename(f) for f in os.listdir(folder_path)
                         if os.path.isfile(os.path.join(folder_path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
                self.processed_files_dict[folder] = files
        self.refresh_category_list()
        self.refresh_processed_list()
        # 更新总数
        total_files = sum(len(files) for files in self.processed_files_dict.values())
        self.processed_frame.configure(text=f"已处理图片 ({self.format_number(total_files)}张)")

    def delete_processed_file(self):
        if not self.source_dir:
            self.status_var.set("运行状态：未找到原始目录！")
            return
        output_dir = os.path.join(os.path.dirname(self.source_dir), '输出图片')
        if not os.path.exists(output_dir):
            self.status_var.set("运行状态：输出图片目录不存在！")
            return
        confirm = messagebox.askyesno("确认删除", f"确定要删除整个输出图片目录吗？\n{output_dir}")
        if not confirm:
            return
        try:
            shutil.rmtree(output_dir)
            self.status_var.set(f"运行状态：已删除 {output_dir}")
            self.processed_files_dict.clear()
            self.refresh_category_list()
            self.refresh_processed_list()
        except Exception as e:
            self.status_var.set(f"运行状态：删除失败: {e}")
            messagebox.showerror("错误", f"删除输出图片目录失败: {e}")

    def copy_processed_file(self):
        if not self.source_dir:
            self.status_var.set("运行状态：未找到原始目录！")
            return
        output_dir = os.path.join(os.path.dirname(self.source_dir), '输出图片')
        if not os.path.exists(output_dir):
            self.status_var.set("运行状态：输出图片目录不存在！")
            return
        # 选择目标目录
        target_dir = filedialog.askdirectory(title="选择复制到的目标文件夹")
        if not target_dir:
            return
        dst = os.path.join(target_dir, os.path.basename(output_dir))
        try:
            if os.path.exists(dst):
                shutil.rmtree(dst)
            shutil.copytree(output_dir, dst)
            self.status_var.set(f"运行状态：已复制到 {dst}")
        except Exception as e:
            self.status_var.set(f"运行状态：复制失败: {e}")
            messagebox.showerror("错误", f"复制输出图片目录失败: {e}")

    def open_processed_file_location(self):
        if not self.source_dir:
            self.status_var.set("运行状态：未找到原始目录！")
            return
        output_dir = os.path.join(os.path.dirname(self.source_dir), '输出图片')
        # 跨平台打开输出图片目录
        try:
            if platform.system() == "Windows":
                os.startfile(output_dir)
            elif platform.system() == "Darwin":
                subprocess.run(["open", output_dir])
            else:
                subprocess.run(["xdg-open", output_dir])
            self.status_var.set(f"运行状态：已打开 {output_dir}")
        except Exception as e:
            self.status_var.set(f"运行状态：打开文件夹失败: {e}")
            messagebox.showerror("错误", f"打开文件夹失败: {e}") 